"""
python p2_solution_pyswarm_bird.py
"""
import json
import numpy as np
import os
import time
from pyswarm import pso

# 读取仓库信息
def load_warehouses(filepath):
    with open(filepath, 'r') as f:
        return json.load(f)

# 读取商品的平均库存量
def load_average_inventory(filepath):
    with open(filepath, 'r') as f:
        return json.load(f)

# 读取商品的平均销量
def load_average_sales(filepath):
    with open(filepath, 'r') as f:
        return json.load(f)

# 读取商品关联度
def load_associations(filepath):
    with open(filepath, 'r') as f:
        return json.load(f)

# 读取初始解
def load_initial_solution(filepath):
    with open(filepath, 'r') as f:
        return json.load(f)

# 适应度函数计算
def fitness_function(solution, warehouses, avg_inventory, avg_sales, associations):
    total_cost = 0
    total_inventory_utilization = 0
    total_sales_utilization = 0
    total_association_sum = 0
    
    # 仓库的数据
    warehouse_data = {warehouse['warehouse_id']: warehouse for warehouse in warehouses}

    # 遍历每个商品
    for category_id, warehouse_index in enumerate(solution):
        warehouse_id = list(warehouse_data.keys())[int(warehouse_index)]
        warehouse = warehouse_data[warehouse_id]

        # 当前商品的平均库存和销量
        inventory = avg_inventory[category_id]['average_inventory']
        sales = avg_sales[category_id]['average_sales']

        # 计算成本
        if inventory > 0:
            total_cost += warehouse['daily_cost']
        
        # 计算库存利用率
        total_inventory_utilization += inventory / warehouse['max_inventory']

        # 计算出货量利用率
        total_sales_utilization += sales / warehouse['max_sales']

        # 计算关联度
        for i in range(len(solution)):
            if solution[i] == warehouse_index:
                for j in range(i + 1, len(solution)):
                    if solution[j] == warehouse_index:
                        # 计算关联度
                        assoc = next((item['association'] for item in associations 
                                      if (item['category1'] == str(category_id) and item['category2'] == str(j)) or 
                                         (item['category1'] == str(j) and item['category2'] == str(category_id))), 0)
                        total_association_sum += assoc

    fitness_value = total_cost * total_inventory_utilization * total_sales_utilization * total_association_sum
    
    # 返回 t1, t2, t3, t4 和适应度值
    return total_cost, total_inventory_utilization, total_sales_utilization, total_association_sum, fitness_value

# 主程序
def main():
    # 路径设置
    warehouse_file = '../fujian/fujian3/origin_data/warehouse.json'
    avg_inventory_file = '../fujian/fujian3/data_from_p1/all_average_inventory.json'
    avg_sales_file = '../fujian/fujian3/data_from_p1/all_average_sales.json'
    associations_file = '../fujian/fujian4/origin_data.json'
    initial_solution_file = '../fujian/p2/simple_solution_handwrite/ans.json'

    # 加载数据
    warehouses = load_warehouses(warehouse_file)
    avg_inventory = load_average_inventory(avg_inventory_file)
    avg_sales = load_average_sales(avg_sales_file)
    associations = load_associations(associations_file)
    initial_solution = load_initial_solution(initial_solution_file)

    # 初始化参数
    num_categories = len(avg_inventory)
    num_warehouses = len(warehouses)

    # 生成初始鸟群位置
    positions = []
    for i in range(num_categories):
        if i < len(initial_solution):  # 确保不超出初始解的范围
            key = str(i)
            positions.append(initial_solution[i].get(key, np.random.randint(0, num_warehouses)))
        else:
            positions.append(np.random.randint(0, num_warehouses))

    # 边界设置
    lb = [0] * num_categories  # 下界
    ub = [num_warehouses - 1] * num_categories  # 上界

    # 初始化速度
    velocities = np.random.rand(num_categories) * (np.array(ub) - np.array(lb)) / 10  # Random initial velocities

    # 保存最佳解
    best_solution_log_file = f'../fujian/p2/bird/best_ever/best_solution.json'
    best_t_log_file = f'../fujian/p2/bird/best_ever/best_t_{time.strftime("%Y%m%d-%H%M%S")}.log'

    # 最佳解和适应度初始化
    global_best_solution = np.copy(positions)
    _, _, _, _, global_best_fitness = fitness_function(global_best_solution, warehouses, avg_inventory, avg_sales, associations)

    # 运行PSO
    maxiter = 1000
    swarmsize = 50  # Number of particles
    for iter in range(maxiter):
        for i in range(swarmsize):
            # Evaluate fitness
            t1, t2, t3, t4, fitness = fitness_function(positions, warehouses, avg_inventory, avg_sales, associations)

            # Update global best
            if fitness < global_best_fitness:
                global_best_fitness = fitness
                global_best_solution = np.copy(positions)

                # Print the objectives along with the best parameters
                print(f"Iteration {iter}: Best parameters - {global_best_solution}, t1: {t1}, t2: {t2}, t3: {t3}, t4: {t4}")

            # Update velocities and positions
            r1 = np.random.rand(num_categories)
            r2 = np.random.rand(num_categories)
            velocities = 0.5 * velocities + r1 * (global_best_solution - positions) + r2 * (positions - np.array(positions))
            positions += velocities

            # Apply boundaries
            positions = np.clip(positions, lb, ub)

        # Print best parameters every 100 iterations
        if iter % 100 == 0:
            print(f"Iteration {iter}: Best parameters - {global_best_solution}")

    # 保存最佳解
    with open(best_solution_log_file, 'w') as f:
        json.dump([{str(i): int(global_best_solution[i])} for i in range(num_categories)], f)

    # 输出当前的适应度到日志文件
    with open(best_t_log_file, 'a') as log:
        log.write(f"{time.time()}: {global_best_fitness}\n")

    print(f"Best solution saved to {best_solution_log_file}")

if __name__ == '__main__':
    main()
